﻿using System;
using System.Collections.Generic;
using System.Net;
using System.Runtime.InteropServices;
using Ast.Core.WinApi.IpHlpApi;
using Ast.Framework;

namespace Ast.Core.Services
{
    public interface INetworkService : IDependency
    {
        IEnumerable<TcpRecord> GetTcpConnections();
    }

    public class NetworkService : INetworkService
    {
        public const int AF_INET = 2; // IPv4
        public const int AF_INET6 = 23; // IPv6

        public IEnumerable<TcpRecord> GetTcpConnections()
        {
            var buffSize = 0;
            var status = IpHelpNative.GetExtendedTcpTable(IntPtr.Zero, ref buffSize, true, AF_INET, IpHelpConstant.TcpTableClass.TCP_TABLE_OWNER_PID_ALL, 0);
            if (status != 122) yield break;

            var buffTable = Marshal.AllocHGlobal(buffSize);
            status = IpHelpNative.GetExtendedTcpTable(buffTable, ref buffSize, true, AF_INET, IpHelpConstant.TcpTableClass.TCP_TABLE_OWNER_PID_ALL, 0);
            if (status != 0) yield break;

            try
            {
                var mibTcpTable = Marshal.PtrToStructure<MibTcptable>(buffTable);
                if (mibTcpTable.NumEntries == 0) yield break;

                var rowIndex = buffTable + Marshal.SizeOf<MibTcptable>();
                for (var i = 0; i < mibTcpTable.NumEntries; i++)
                {
                    var mibTcpRow = Marshal.PtrToStructure<MibTcprow>(rowIndex);
                    yield return new TcpRecord(
                        new IPAddress(mibTcpRow.LocalAddress),
                        new IPAddress(mibTcpRow.RemoteAddress),
                        (ushort) IPAddress.NetworkToHostOrder(mibTcpRow.LocalPort),
                        (ushort) IPAddress.NetworkToHostOrder(mibTcpRow.RemotePort),
                        mibTcpRow.ProcessId,
                        mibTcpRow.State
                        );

                    rowIndex += Marshal.SizeOf<MibTcprow>();
                }
            }
            finally
            {
                Marshal.FreeHGlobal(buffTable);
            }
        }
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto, Pack = 1)]
    public struct MibTcptable
    {
        public uint NumEntries;
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto, Pack = 1)]
    public struct MibTcprow
    {
        public uint State;

        public uint LocalAddress;

        public short LocalPort;
        private short _;

        public uint RemoteAddress;

        public short RemotePort;
        private short __;

        public int ProcessId;
    }

    public class TcpRecord
    {
        public IPAddress LocalAddress { get; set; }
        public IPAddress RemoteAddress { get; set; }
        public ushort LocalPort { get; set; }
        public ushort RemotePort { get; set; }
        public int ProcessId { get; set; }
        public uint State { get; set; }

        public TcpRecord(IPAddress localIp, IPAddress remoteIp, ushort localPort, ushort remotePort, int processId, uint state)
        {
            LocalAddress = localIp;
            RemoteAddress = remoteIp;
            LocalPort = localPort;
            RemotePort = remotePort;
            ProcessId = processId;
            State = state;
        }
    }
}